#include "nm_inlinehook.h"

VOID __stdcall kinlinehook::DpcLock(struct _KDPC *Dpc, PVOID DeferredContext, PVOID SystemArgument1, PVOID SystemArgument2)
{
	__asm cli;

	InterlockedIncrement(&((PDPCLOCK_CONTEXT)DeferredContext)->Locks);
	do {
		__asm pause;

	} while (!((PDPCLOCK_CONTEXT)DeferredContext)->Release);

	InterlockedDecrement(&((PDPCLOCK_CONTEXT)DeferredContext)->Locks);
	__asm sti;
}


BOOLEAN kinlinehook::AutoHook(PVOID Dst, PVOID Src, PVOID *Proxy)
{
	KIRQL OldIrql;
	PRKDPC Dpc;
	DPCLOCK_CONTEXT DpcLockContext;
	CCHAR Processor;
	const ULONG MicroSeconds = 1000 * 50;
	const ULONG Tag = 'nmlk';
	
	UCHAR JmpTmp[5] = {0xe9, 0x90, 0x90, 0x90, 0x90};
	ULONG CodeLength = 0;
	PUCHAR OpCodePtr;
	ULONG TmpJmpAddr = 0;
	ULONG TmpJmpOff = 0;

	//
	// Prepare hook.
	//
	if (Dst == NULL || Src == NULL || Proxy == NULL) {

		return FALSE;
	}

	//
	// Calculate code length.
	//

	while (CodeLength < 5) {

		CodeLength += SizeOfCode((PUCHAR)Dst + CodeLength, &OpCodePtr);

		if (*OpCodePtr == 0xe8 || *OpCodePtr == 0xe9) {
			TmpJmpAddr = (ULONG)(OpCodePtr + 5 + *(ULONG *)(OpCodePtr + 1));
			TmpJmpOff = CodeLength - 5;
		}

		//
		// Not support some code.
		//

		if (*OpCodePtr == 0xeb) {

			//
			// I will not support some opcode, such as 0xeb and so on.
			//
			return FALSE;
		}
	}

	//
	// Prepare proxy routine addr. And calculate jump address.
	//

	*Proxy = ExAllocatePoolWithTag(NonPagedPool, CodeLength + 5, Tag);
	RtlCopyMemory(*Proxy, Dst, CodeLength);
	*((PUCHAR)*Proxy + CodeLength) = 0xe9;
	*(ULONG *)((PUCHAR)*Proxy + CodeLength + 1) = (ULONG)Dst - ((ULONG)*Proxy + 5);

	if (TmpJmpAddr != 0) {

		*(ULONG *)((PUCHAR)*Proxy + TmpJmpOff + 1) = TmpJmpAddr - ((ULONG)*Proxy + TmpJmpOff + 5);
	}

	*(ULONG *)&JmpTmp[1] = (ULONG)Src - ((ULONG)Dst + 5);


	//
	// Raise Irql
	//
	OldIrql = KeRaiseIrqlToDpcLevel();

	if (KeNumberProcessors > 1) {

		//
		// Let current thread work on the processor 0.
		//
		KeSetAffinityThread(PsGetCurrentThread(), (KAFFINITY)1);
		if (KeGetCurrentProcessorNumber() != 0) {

			ExFreePoolWithTag(*Proxy, Tag);
			*Proxy = NULL;
			return FALSE;
		}

		DpcLockContext.Dpcs = (PKDPC)ExAllocatePoolWithTag(NonPagedPool, (KeNumberProcessors - 1) * sizeof(KDPC), Tag);

		for (Processor = 1; Processor < KeNumberProcessors; Processor++) {

			Dpc = &DpcLockContext.Dpcs[Processor];
			KeInitializeDpc(Dpc, kinlinehook::DpcLock, &DpcLockContext);
			KeSetImportanceDpc(Dpc, HighImportance);
			KeSetTargetProcessorDpc(Dpc, Processor);
			KeInsertQueueDpc(Dpc, NULL, NULL);

		}

		while (DpcLockContext.Locks != (KeNumberProcessors - 1)) {

			KeStallExecutionProcessor(MicroSeconds);
		}
		
		kinlinehook::WPOFF();
		
		RtlCopyMemory(Dst, JmpTmp, 5);

		kinlinehook::WPOFF();


		DpcLockContext.Release = TRUE;
		while (DpcLockContext.Locks != 0) {

			KeStallExecutionProcessor(MicroSeconds);
		}

		ExFreePoolWithTag(DpcLockContext.Dpcs, Tag);

	} 
	else {

		kinlinehook::WPOFF();
		
		RtlCopyMemory(Dst, JmpTmp, 5);

		kinlinehook::WPOFF();
	}

	KeLowerIrql(OldIrql);

	return TRUE;
}

BOOLEAN kinlinehook::Hook(HOOKPROC HookProc, PVOID Context)
{
	KIRQL OldIrql;
	PRKDPC Dpc;
	DPCLOCK_CONTEXT DpcLockContext;
	CCHAR Processor;
	const ULONG MicroSeconds = 1000 * 50;
	const ULONG Tag = 'nmlk';
	BOOLEAN Ret = FALSE;

	//
	// Raise Irql
	//
	OldIrql = KeRaiseIrqlToDpcLevel();

	if (KeNumberProcessors > 1) {

		//
		// Let current thread work on the processor 0.
		//
		KeSetAffinityThread(PsGetCurrentThread(), (KAFFINITY)1);
		if (KeGetCurrentProcessorNumber() != 0) {

			return FALSE;
		}

		DpcLockContext.Dpcs = (PKDPC)ExAllocatePoolWithTag(NonPagedPool, (KeNumberProcessors - 1) * sizeof(KDPC), Tag);

		for (Processor = 1; Processor < KeNumberProcessors; Processor++) {

			Dpc = &DpcLockContext.Dpcs[Processor];
			KeInitializeDpc(Dpc, kinlinehook::DpcLock, &DpcLockContext);
			KeSetImportanceDpc(Dpc, HighImportance);
			KeSetTargetProcessorDpc(Dpc, Processor);
			KeInsertQueueDpc(Dpc, NULL, NULL);

		}

		while (DpcLockContext.Locks != (KeNumberProcessors - 1)) {

			KeStallExecutionProcessor(MicroSeconds);
		}

		kinlinehook::WPOFF();
		
		if (ARGUMENT_PRESENT(HookProc)) {
			Ret = HookProc(Context);
		}
		
		kinlinehook::WPOFF();


		DpcLockContext.Release = TRUE;
		while (DpcLockContext.Locks != 0) {

			KeStallExecutionProcessor(MicroSeconds);
		}

		ExFreePoolWithTag(DpcLockContext.Dpcs, Tag);

	} 
	else {

		kinlinehook::WPOFF();

		if (ARGUMENT_PRESENT(HookProc)) {
			Ret = HookProc(Context);
		}

		kinlinehook::WPOFF();
	}

	KeLowerIrql(OldIrql);

	return Ret;
}